
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "testDSAKeys.h"
#include "../utils/StringUtils.h"

// add a test for pipe exhaustion - leave stuff in the buffer and call pipeFinal
// add signature verification/checking with no key and see what happens

void tests::testDSAKeys(apr_pool_t *mp)
{
   // Can generate a new DSA key pair as:
   DSAKeyPairGenerator dsakpgen(DSA_BITS, "DSAkey");
   DSAPublicKey *pubkey = dsakpgen.popPublicKey();
   DSAPrivateKey *prvkey = dsakpgen.popPrivateKey();
   DSAKeyPairGenerator dsakpgen2(DSA_BITS, "DSAkey2");
   DSAPublicKey *bpubkey = dsakpgen2.popPublicKey();
   DSAPrivateKey *bprvkey = dsakpgen2.popPrivateKey();
   
   int iters = 10;
   testDSAPasswordCryption(pubkey, prvkey, iters);
   testDSABlockSigning(bpubkey, bprvkey, pubkey, prvkey, iters, 10000);
   testDSAPipeSigning(bpubkey, bprvkey, pubkey, prvkey, iters, 500, 10000, mp);
   
   delete pubkey;
   delete prvkey;
   delete bpubkey;
   delete bprvkey;
   ERR_remove_state(0);
}

void tests::testDSAPasswordCryption(DSAPublicKey *pubkey, DSAPrivateKey *prvkey, int iterations)
{
   for (int i = 0; i < iterations; i++)
   {
      char password_buffer[256];
      char *exporttext = NULL;
      sprintf(password_buffer, "password%d", rand());
      pubkey->exportKeyToASCIIWithPassword(&exporttext, password_buffer);
      DSAPublicKey *pubkey2 = new DSAPublicKey("newDSA");
      pubkey2->importKeyFromASCIIWithPassword(exporttext, password_buffer, "newDSA");
      int compare_key = pubkey2->compareKeyEqual(pubkey);
      if (compare_key)
      {
         printf("DSA PasswordCryption PublicKey test %d passed - password '%s'\n", i, password_buffer);
      } else
      {
         errprint("failed: encrypted key was '%s'", exporttext);
         char *data1 = NULL, *data2 = NULL;
         pubkey->exportKeyToASCII(&data1);
         pubkey2->exportKeyToASCII(&data2);
         errprint("original key base64 '%s'", data1);
         errprint("final key base64 '%s'", data2);
         free(data1); free(data2);
      }
      free(exporttext);

      compare_key = prvkey->compareKeyBelongs(pubkey);
      if (compare_key)
      {
         printf("DSA Matches PrivateKey test %d passed - password '%s'\n", i, password_buffer);
      } else
      {
         errprint("failed: key was '%s'", exporttext);
      }
      prvkey->exportKeyToASCIIWithPassword(&exporttext, password_buffer);
      DSAPrivateKey *prvkey2 = new DSAPrivateKey("newDSA");
      prvkey2->importKeyFromASCIIWithPassword(exporttext, password_buffer, "newDSA");
      compare_key = prvkey2->compareKeyBelongs(pubkey);
      if (compare_key)
      {
         printf("DSA PasswordCryption PrivateKey test %d passed - password '%s'\n", i, password_buffer);
      } else
      {
         errprint("failed: encrypted key was '%s'", exporttext);
         char *data1 = NULL, *data2 = NULL;
         prvkey->exportKeyToASCII(&data1);
         prvkey2->exportKeyToASCII(&data2);
         errprint("original key base64 '%s'", data1);
         errprint("final key base64 '%s'", data2);
         free(data1); free(data2);
      }
      free(exporttext);
      delete pubkey2;
      delete prvkey2;
   }
}

void tests::testDSABlockSigning(DSAPublicKey *bpubkey, DSAPrivateKey *bprvkey, DSAPublicKey *pubkey, DSAPrivateKey *prvkey, int iterations, int max_length)
{
   for (int i = 0; i < iterations; i++)
   {
      // Random length much more than a single key
      int randlen = rand() % max_length;
      // Create a block of (pseudo-)random data
      char *pseudorandom = (char *)malloc(sizeof(char)*randlen);
      for (int j = 0; j < randlen; j++)
         pseudorandom[j] = (rand() & 0xFF0) >> 8;
      // Sign the random data
      apr_size_t new_coded_size = randlen;
      apr_size_t new_plain_size = 0;
      DSASigning *dsa_sgn = new DSASigning(*prvkey);
      char *encrypted_data = dsa_sgn->processDataBlock(pseudorandom, randlen, &new_coded_size);
      delete dsa_sgn;
      DSAChecking *dsa_chk = NULL;
      int totest = rand() % 3;
      if (totest)
         dsa_chk = new DSAChecking(*pubkey);
      else
         dsa_chk = new DSAChecking(*bpubkey);
      if (totest == 2)
         dsa_chk->removeKey();
      char *new_plain_data = dsa_chk->processDataBlock(encrypted_data, new_coded_size, &new_plain_size);
      int verify_data = dsa_chk->processWasSecure();
      delete dsa_chk;
      if (new_plain_size != randlen)
      {
         errprint("error: original data size was %" APR_SIZE_T_FMT ", new data size is %" APR_SIZE_T_FMT, (apr_size_t)randlen, new_plain_size);
      } else if (memcmp(pseudorandom, new_plain_data, sizeof(char)*randlen) != 0)
      {
         errprint("error: original data does not match new data");
         //printf("old data:\n");
         //StringUtils::printBinary(pseudorandom, randlen);
         //printf("new data:\n");
         //StringUtils::printBinary(new_plain_data, randlen);
      } else
      {
         printf("DSA BlockSigning(issame) test %d passed - data length %d\n", i, randlen);
         const char *testtype = "";
         switch (totest)
         {
            case 0: testtype = "bad key"; break;
            case 1: testtype = "correct key"; break;
            case 2: testtype = "missing key"; break;
         }
         printf("DSA BlockCheck %s (%s: needed a %d, got a %d)\n",
            ((totest & 1) == verify_data) ? "passed" : "failed", testtype, totest & 1, verify_data);
      }
      free(pseudorandom);
      free(encrypted_data);
      free(new_plain_data);
   }
}

void tests::testDSAPipeSigning(DSAPublicKey *bpubkey, DSAPrivateKey *bprvkey, DSAPublicKey *pubkey, DSAPrivateKey *prvkey, int iterations, int max_block, int max_length, apr_pool_t *mp)
{
   for (int i = 0; i < iterations; i++)
   {
      // Random length much more than a single key
      int randlen = rand() % max_length;
      // Create a block of (pseudo-)random data
      char *pseudorandom = (char *)malloc(sizeof(char)*randlen);
      for (int j = 0; j < randlen; j++)
         pseudorandom[j] = (rand() & 0xFF0) >> 8;
      // Encrypt the random data
      apr_size_t new_coded_size = randlen;
      APRBucketBrigade p_input(mp);
      APRBucketBrigade coded(mp);
      DSASigning *dsa_sgn = new DSASigning(p_input, coded, *prvkey);
      dsa_sgn->processPipeBegin();
      int nidx = 0;
      while (nidx < randlen)
      {
         int randblocklen = rand() % max_block;
         randblocklen = ((randblocklen + nidx) > randlen) ? (randlen - nidx) : randblocklen;
         int randblocklen_added = p_input.writeToStream(&(pseudorandom[nidx]), randblocklen);
         nidx += randblocklen_added;
         dsa_sgn->processPipeStep();
      }
      dsa_sgn->processPipeEnd();
      delete dsa_sgn;
      APRBucketBrigade coded2(mp);
      APRBucketBrigade p_output(mp);
      DSAChecking *dsa_chk = NULL;
      int totest = rand() % 3;
      if (totest)
         dsa_chk = new DSAChecking(coded2, p_output, *pubkey);
      else
         dsa_chk = new DSAChecking(coded2, p_output, *bpubkey);
      if (totest == 2)
         dsa_chk->removeKey();
      dsa_chk->processPipeBegin();
      nidx = 0;
      int total_length = coded.realLength();
      while (nidx < total_length)
      {
         int randblocklen = rand() % max_block;
         randblocklen = ((randblocklen + nidx) > total_length) ? (total_length - nidx) : randblocklen;
         char *working_buffer = (char *)malloc(sizeof(char)*randblocklen);
         int block_read = coded.readFromStream(working_buffer, randblocklen);
         if (block_read != randblocklen)
         {
            errprint("read wrong amount from bucket (expected %d, got %d)", randblocklen, block_read);
         } else
            coded.removeFromStream(randblocklen);
         int block_write = coded2.writeToStream(working_buffer, randblocklen);
         if (block_write != randblocklen)
         {
            errprint("wrote wrong amount to bucket (expected %d, wrote %d)", randblocklen, block_write);
         }
         nidx += randblocklen;
         dsa_chk->processPipeStep();
         free(working_buffer);
      }
      dsa_chk->processPipeEnd();
      int verify_data = dsa_chk->processWasSecure();
      delete dsa_chk;
      size_t new_plain_size = 0;
      char *new_plain_data = NULL;
      p_output.bucketBrigadeToBuffer(&new_plain_data, &new_plain_size);
      if (new_plain_size != randlen)
      {
         errprint("error: original data size was %" APR_SIZE_T_FMT ", new data size is %" APR_SIZE_T_FMT, (apr_size_t)randlen, new_plain_size);
         printf("old data:\n");
         StringUtils::printBinary(pseudorandom, randlen);
         printf("new data:\n");
         StringUtils::printBinary(new_plain_data, randlen);
      } else if (memcmp(pseudorandom, new_plain_data, sizeof(char)*randlen) != 0)
      {
         errprint("error: original data does not match new data");
         //printf("old data:\n");
         //StringUtils::printBinary(pseudorandom, randlen);
         //printf("new data:\n");
         //StringUtils::printBinary(new_plain_data, randlen);
      } else
      {
         printf("DSA PipeSigning(issame) test %d passed - data length %d\n", i, randlen);
         const char *testtype = "";
         switch (totest)
         {
            case 0: testtype = "bad key"; break;
            case 1: testtype = "correct key"; break;
            case 2: testtype = "missing key"; break;
         }
         printf("DSA PipeCheck %s (%s: needed a %d, got a %d)\n",
            ((totest & 1) == verify_data) ? "passed" : "failed", testtype, totest & 1, verify_data);
      }
      free(pseudorandom);
      free(new_plain_data);
   }
}
